import React, { useEffect, useRef } from 'react';

interface AttentionHeatmapProps {
  attentionWeights: number[][];
  sequenceLength: number;
}

const AttentionHeatmap: React.FC<AttentionHeatmapProps> = ({ 
  attentionWeights,
  sequenceLength 
}) => {
  const canvasRef = useRef<HTMLCanvasElement>(null);
  
  useEffect(() => {
    if (!canvasRef.current || !attentionWeights.length) return;
    
    const canvas = canvasRef.current;
    const ctx = canvas.getContext('2d');
    if (!ctx) return;
    
    // Set canvas dimensions
    const cellSize = 40;
    const padding = 40;
    const width = sequenceLength * cellSize + padding * 2;
    const height = sequenceLength * cellSize + padding * 2;
    
    canvas.width = width;
    canvas.height = height;
    
    // Clear canvas
    ctx.clearRect(0, 0, width, height);
    
    // Draw heatmap
    for (let i = 0; i < sequenceLength; i++) {
      for (let j = 0; j < sequenceLength; j++) {
        const value = attentionWeights[i][j];
        
        // Calculate color based on value (blue gradient)
        const intensity = Math.floor(255 * value);
        const red = 0;
        const green = Math.floor(intensity * 0.5);
        const blue = intensity;
        
        ctx.fillStyle = `rgb(${red}, ${green}, ${blue})`;
        ctx.fillRect(
          j * cellSize + padding, 
          i * cellSize + padding, 
          cellSize, 
          cellSize
        );
        
        // Add value text
        ctx.fillStyle = value > 0.5 ? 'white' : 'black';
        ctx.font = '10px sans-serif';
        ctx.textAlign = 'center';
        ctx.textBaseline = 'middle';
        ctx.fillText(
          value.toFixed(2), 
          j * cellSize + padding + cellSize / 2, 
          i * cellSize + padding + cellSize / 2
        );
      }
    }
    
    // Draw labels on axes
    ctx.fillStyle = 'black';
    ctx.font = '12px sans-serif';
    ctx.textAlign = 'center';
    
    for (let i = 0; i < sequenceLength; i++) {
      // Draw column labels (position in sequence)
      ctx.fillText(
        `t-${sequenceLength - i}`, 
        i * cellSize + padding + cellSize / 2, 
        padding - 15
      );
      
      // Draw row labels (position in sequence)
      ctx.textAlign = 'right';
      ctx.fillText(
        `t-${sequenceLength - i}`, 
        padding - 10, 
        i * cellSize + padding + cellSize / 2
      );
      ctx.textAlign = 'center';
    }
    
    // Draw axes titles
    ctx.font = '14px sans-serif';
    ctx.fillText('Position in Sequence', width / 2, padding - 25);
    
    // Rotate text for y-axis
    ctx.save();
    ctx.translate(padding - 25, height / 2);
    ctx.rotate(-Math.PI / 2);
    ctx.fillText('Attention Focus', 0, 0);
    ctx.restore();
    
  }, [attentionWeights, sequenceLength]);
  
  if (!attentionWeights.length) {
    return (
      <div className="bg-white rounded-lg border border-gray-200 p-6 text-center shadow-sm">
        <p className="text-gray-500">No attention data available</p>
      </div>
    );
  }
  
  return (
    <div className="bg-white rounded-lg border border-gray-200 p-4 shadow-sm">
      <h3 className="text-lg font-medium text-gray-800 mb-4">Attention Heatmap</h3>
      <div className="overflow-auto flex justify-center">
        <canvas 
          ref={canvasRef}
          className="bg-white"
        />
      </div>
      <p className="text-sm text-gray-500 mt-2 text-center">
        The heatmap shows how much attention the model pays to each position when making predictions.
        Darker colors indicate stronger attention.
      </p>
    </div>
  );
};

export default AttentionHeatmap;